from starlette.middleware.base import BaseHTTPMiddleware
from fastapi.responses import JSONResponse
from models import SessionLocal, SendWhitelist, get_now_local
import ipaddress

class SendWhitelistMiddleware(BaseHTTPMiddleware):
    """
    发送白名单中间件：仅拦截 POST /send
    允许来源：命中白名单（ip/cidr/domain）或默认本地段（192.168.*）
    """
    def __init__(self, app):
        super().__init__(app)
        self._cache = {"data": None, "expires": None}

    async def dispatch(self, request, call_next):
        if request.method == "POST" and request.url.path == "/send":
            now = get_now_local()
            # 缓存 60 秒
            need_reload = True
            if self._cache["expires"] and self._cache["expires"] > now and self._cache["data"] is not None:
                need_reload = False
            if need_reload:
                try:
                    with SessionLocal() as s:
                        rows = s.query(SendWhitelist).filter(SendWhitelist.enabled == 1).all()
                        self._cache["data"] = [(r.type, r.value) for r in rows]
                        from datetime import timedelta
                        self._cache["expires"] = now + timedelta(seconds=60)
                except Exception:
                    self._cache["data"] = []
                    self._cache["expires"] = now

            src_ip = (request.headers.get("x-forwarded-for") or "").split(",")[0].strip() or request.headers.get("x-real-ip") or (request.client.host if request.client else None)
            host = request.headers.get("host") or request.url.hostname

            allowed = False
            wl = self._cache["data"] or []
            if src_ip and src_ip.startswith("192.168."):
                allowed = True
            try:
                # IP 精确与 CIDR
                if src_ip:
                    for t, v in wl:
                        if t == "ip" and src_ip == v:
                            allowed = True
                            break
                        if t == "cidr":
                            try:
                                net = ipaddress.ip_network(v, strict=False)
                                ipobj = ipaddress.ip_address(src_ip)
                                if ipobj in net:
                                    allowed = True
                                    break
                            except Exception:
                                continue
                # 域名
                if not allowed and host:
                    for t, v in wl:
                        if t == "domain" and host.split(":")[0] == v:
                            allowed = True
                            break
            except Exception:
                allowed = False

            if not allowed:
                return JSONResponse({"code": 403, "msg": "forbidden: not in whitelist"}, status_code=403)

        return await call_next(request)

